#include <common.h>
#include <elf.h>

FILE *ftrace_fp = NULL;
Elf32_Off symtab = 0;
Elf32_Off strtab = 0;
uint32_t symtab_size = 0;
uint32_t strtab_size = 0;
int call_num = 0;

typedef struct Stack {
    uint32_t pc;
    uint32_t dnpc;
    struct Stack *next;
} Stack;

Stack *stack = NULL;

void stack_pop() {
  Stack *sp = stack;
  stack = stack->next;

  free(sp);
}

void stack_push(uint32_t pc, uint32_t dnpc) {
  Stack *sp = malloc(sizeof(Stack));
  Assert(sp, "Ftrace: stack overflow.\n");

  sp->pc = pc;
  sp->dnpc = dnpc;
  sp->next = stack;
  stack = sp;
}

void stack_free() {
  Stack *tmp = NULL;
  while (stack != NULL) {
    tmp = stack;
    stack = stack->next;
    free(tmp);
  }
}

int stack_find(uint32_t pc, uint32_t *dnpc) {
  if (stack != NULL) {
    if (stack->pc == pc) {
      *dnpc = stack->dnpc;
      return 1;
    }
  }

  return 0;
}

void ftrace_end() {
  fclose(ftrace_fp);
  stack_free();
}

void init_ftrace(const char *ftrace_file) {
  if (ftrace_file != NULL) {
    FILE *fp = fopen(ftrace_file, "rb");
    Assert(fp, "Can not open '%s'", ftrace_file);
    ftrace_fp = fp;

    Elf32_Ehdr elf_header;
    fread(&elf_header, 1, sizeof(Elf32_Ehdr), fp);

    // get section header address
    Elf32_Off e_shoff = elf_header.e_shoff;
    if (e_shoff == 0) printf("No section header\n");

    // get section header string table index
    unsigned int shdr_size = sizeof(Elf32_Shdr);
    Elf32_Shdr shdr;
    uint32_t shstrtab = 0;
    uint16_t e_shstrndx = elf_header.e_shstrndx;
    if (e_shstrndx == SHN_UNDEF) printf("No section name string table\n");
    else if (e_shstrndx >= SHN_LORESERVE) {
      // real index is sh_link member of the initial entry in section header table
      if (fseek(fp, e_shoff, SEEK_SET) != 0) printf("ERROR\n");

      fread(&shdr, 1, shdr_size, fp);
      shstrtab = shdr.sh_link;
    }
    else shstrtab = e_shstrndx;

    // get shstrtab offset
    fseek(fp, e_shoff + shdr_size * shstrtab, SEEK_SET);
    fread(&shdr, 1, shdr_size, fp);
    Elf32_Off shstrtab_off = shdr.sh_offset;

    // get .symtab and .strtab address
    int count = 0;
    char buf[8];
    uint16_t e_shnum = elf_header.e_shnum;
    for (int i = 0; i < e_shnum; i++) {
      fseek(fp, e_shoff + i * shdr_size, SEEK_SET);
      fread(&shdr, 1, shdr_size, fp);

      fseek(fp, shstrtab_off + shdr.sh_name, SEEK_SET);
      fgets(buf, 8, fp);

      if (strcmp(buf, ".symtab") == 0) {
        symtab = shdr.sh_offset;
        symtab_size = shdr.sh_size;
        count++;
      }
      else if (strcmp(buf, ".strtab") == 0) {
        strtab = shdr.sh_offset;
        strtab_size = shdr.sh_size;
        count++;
      }
      else if (count == 2) break;
    }

#ifndef CONFIG_FTRACE
  ftrace_end();
#endif
  }

  // printf(".symtab:      0x%x\n", symtab);
  // printf(".symtab_size: 0x%x\n", symtab_size);
  // printf(".strtab:      0x%x\n", strtab);
  // printf(".strtab_size: 0x%x\n", strtab_size);
}

void ftrace_find_fn(uint32_t addr, char *name) {
  // find function name
  int sym_size = sizeof(Elf32_Sym);
  int sym_num = symtab_size / sym_size;
  Elf32_Sym sym;
  
  for (int i = 0; i < sym_num; i++) {
    fseek(ftrace_fp, symtab + i * sym_size, SEEK_SET);
    fread(&sym, 1, sym_size, ftrace_fp);

    if(sym.st_value == addr && ELF32_ST_TYPE(sym.st_info) == STT_FUNC) {
      fseek(ftrace_fp, strtab + sym.st_name, SEEK_SET);
      fgets(name, sizeof(name), ftrace_fp);

      break;
    }
  }
}

void ftrace_trace(uint32_t pc, uint32_t snpc, uint32_t dnpc) {
  if (snpc != dnpc) {
    // encountered jump instruction
    char call_name[40] = {'\0'};
    char ret_name[40] = {'\0'};
    uint32_t ret_dnpc = 0;

    ftrace_find_fn(dnpc, call_name);
    if (stack_find(dnpc - 4, &ret_dnpc)) {
      ftrace_find_fn(ret_dnpc, ret_name);
    }

    // display
    if (strcmp(call_name, "") != 0) {
      printf("0x%x: ", pc);
      for (int i = 0; i < call_num; i++) {
        printf("  ");
      }
      printf("call [%s@0x%08x]\n", call_name, dnpc);
      call_num++;

      stack_push(pc, dnpc);
    }
    else {
      if(strcmp(ret_name, "") != 0) {
        call_num--;
        printf("0x%x: ", pc);
        for (int i = 0; i < call_num; i++) {
          printf("  ");
        }
        printf("ret  [%s]\n", ret_name);

        stack_pop();
      }
    }
  }
}

